import itertools
from typing import List, Optional, Tuple

from openai import OpenAI
from openai.types.chat import ChatCompletion, ChatCompletionContentPartParam, ChatCompletionMessageParam

from aicenter.models import AIIdentity, AIModel, AIRole
from .._define import (
    AIChatClient, AIChatContent, AIChatMessage, AIChatRequest, AIError, AITokenUsage, ContentType,
    MessageRole,
)

# 默认参数集
# 用于定义一些默认传递给接口的参数
chat_default_params = {
    # 核采样
    # 温度采样的另一种方法，称为核采样，模型会考虑概率质量为 top_p 的标记的结果。
    # 因此 0.1 意味着只考虑概率质量最高的 10%的标记。
    # 我们通常建议更改此值或 `temperature`，但不能同时更改。
    'top_p': 1.0,
    # 采样温度
    # 使用的采样温度，介于 0 和 2 之间。
    # 0.8 等较高值会使输出更加随机，而 0.2 等较低值则会使输出更加集中和确定。
    # 我们通常建议改变该值或 `top_p`，但不要同时改变这两个值。
    'temperature': 1.0,
    # 逻辑偏差
    # 修改指定标记在完成时出现的可能性。
    # 接受一个 JSON 对象，该对象将标记（由标记化器中的标记 ID 指定）映射到相关的偏置值（从 -100 到 100）。
    # 在数学上，偏差会在采样前添加到模型生成的对数中。
    # 具体效果因模型而异，但 -1 和 1 之间的值会降低或增加选择的可能性；
    # -100 或 100 这样的值会导致禁止或完全选择相关标记。
    'logit_bias': None,
    # 存在惩罚
    # 介于 -2.0 和 2.0 之间的数值。
    # 正值会根据新标记是否出现在文本中对其进行惩罚，从而增加模型谈论新话题的可能性。
    'presence_penalty': 0.0,
    # 频率处罚
    # 介于 -2.0 和 2.0 之间的数值。
    # 正值会根据新标记在文本中的现有频率对其进行惩罚，从而降低模型逐字重复同一行的可能性。
    'frequency_penalty': 0.0,
}
# 固定参数集
# 用于定义一些不允许接口调用方修改的参数
chat_fixed_params = {
    # 响应结果数
    # 为每条输入信息生成多少个聊天完成选项。
    'n': 1,
    # 响应格式
    # 指定模型必须输出的格式的对象。
    # 设置为 `{ "type"："json_object" }` 可启用 JSON 模式，从而保证模型生成的信息是有效的 JSON。
    # 重要：
    #     使用 JSON 模式时，还必须通过系统或用户消息指示模型自己生成 JSON。
    #     如果不这样做，模型可能会生成无休止的空白流，直到生成量达到标记限制，从而导致请求长时间运行，看似 "卡住了"。
    #     还要注意的是，如果 `finish_reason="length"`，消息内容可能会被部分截断，
    #     这表明生成超过了 `max_tokens`，或对话超过了最大上下文长度。
    'response_format': None,
}


class AIContinueChatRequest(AIChatRequest):
    def __init__(self, request: AIChatRequest, last_msg: AIChatMessage):
        super().__init__(
            request.model,
            request.role,
            request.identity,
            request.history,
            request.messages,
            request.params,
        )
        self._last_msg = last_msg

    def messages(self) -> Tuple[AIChatMessage, ...]:
        return *super().messages, self._last_msg


class OpenAIClient(AIChatClient):
    def _initialize(self, params: dict):
        self._openai = OpenAI(
            api_key=self.api_key,
            organization=self.app_id,
        )

    def _register_model(self, model: AIModel) -> bool:
        return True

    def _register_identity(self, model: AIModel, role: AIRole, identity: AIIdentity) -> bool:
        return True

    @staticmethod
    def _extract_end_user(request: AIChatRequest) -> str:
        """提取最终用户"""
        # 从后向前遍历消息
        for message in reversed(request.messages):
            # 若消息不是用户发出的，则跳过
            if message.role != MessageRole.User:
                continue
            # 返回最后一个发出消息的用户的名称
            return message.speaker
        # 若找不到用户发出的消息，则将最后一条消息的发言者视为最终用户
        return request.messages[-1].speaker

    @staticmethod
    def _build_content(role: MessageRole, content: AIChatContent) -> ChatCompletionContentPartParam:
        """构建内容"""
        if role != MessageRole.User and content.type != ContentType.Text:
            raise AIError(f'OpenAI -> 非用户消息仅支持 Text 类型的内容。')
        if content.type == ContentType.Text:
            # 构建文本消息
            return {
                'type': 'text',
                'text': content.content,
            }
        elif content.type == ContentType.Image:
            # 构建图像消息
            return {
                'type': 'image_url',
                'image_url': {
                    'url': content.content,
                },
            }
        raise AIError(f'OpenAI -> 不支持的内容类型：${content.type.name}')

    @staticmethod
    def _build_message(message: AIChatMessage) -> ChatCompletionMessageParam:
        """构建消息"""
        contents = [OpenAIClient._build_content(message.role, content) for content in message.contents]
        if message.role == MessageRole.User:
            return {
                'role': 'user',
                'content': contents,
            }
        elif message.role == MessageRole.System:
            return {
                'role': 'system',
                'content': contents,
            }
        raise AIError(f'OpenAI -> 不支持的消息角色：${message.role.name}')

    @staticmethod
    def _build_messages(request: AIChatRequest) -> List[ChatCompletionMessageParam]:
        """构建消息"""
        return [
            OpenAIClient._build_message(message)
            for message in itertools.chain(request.history, request.messages)
        ]

    @staticmethod
    def _extract_messages(model: AIModel, resp: ChatCompletion) -> AIChatMessage:
        """提取消息"""
        return AIChatMessage(
            MessageRole.AI,
            model.name,
            [
                AIChatContent(
                    ContentType.Text,
                    resp.choices[0].message.content,
                ),
            ],
        )

    @staticmethod
    def _extract_token_usage(resp: ChatCompletion) -> AITokenUsage:
        """提取令牌使用情况"""
        usage = resp.usage
        # 由于 OpenAI 的接口将用户的输入内容也视为提示词，
        # 故此处将提示词令牌使用量视为输入令牌使用量
        return AITokenUsage(
            # 提示词令牌使用量
            0,
            # 输入令牌使用量
            usage.prompt_tokens,
            # 输出令牌使用量
            usage.completion_tokens,
            # 总令牌使用量
            usage.total_tokens,
        )

    @staticmethod
    def _merge_result(
            model: AIModel,
            last_msg: AIChatMessage, next_msg: AIChatMessage,
            last_token_usage: AITokenUsage, next_token_usage: AITokenUsage,
    ) -> (AIChatMessage, AITokenUsage):
        """合并两次对话的结果"""
        msg = AIChatMessage(
            MessageRole.AI,
            model.name,
            [
                AIChatContent(
                    ContentType.Text,
                    last_msg.contents[0].content + next_msg.contents[0].content
                ),
            ],
        )
        token_usage = AITokenUsage(
            last_token_usage.prompt_usage + next_token_usage.prompt_usage,
            last_token_usage.input_usage + next_token_usage.input_usage,
            last_token_usage.output_usage + next_token_usage.output_usage,
            last_token_usage.total_usage + next_token_usage.total_usage,
        )
        return msg, token_usage

    def _try_continue_chat(
            self,
            reason: str,
            request: AIChatRequest,
            last_msg: AIChatMessage,
    ) -> (bool, Optional[AIChatMessage], Optional[AITokenUsage]):
        """继续对话"""
        # 若终止原因为 `超过最大令牌限制`，则自动继续对话
        if reason == 'length':
            # 创建新的 AI 请求
            new_req = AIContinueChatRequest(request, last_msg)
            # 发起对话
            next_messages, next_token_usage = self._chat(new_req)
            # 返回新的对话结果
            return next_messages[-1], next_token_usage
        # 若为其他原因，则不执行继续对话
        return False, None, None

    def _chat(self, request: AIChatRequest) -> (Tuple[AIChatMessage], AITokenUsage):
        # 构建参数集合
        params_array = [chat_fixed_params, request.params, request.model.model_params, chat_default_params]
        # 调用接口
        resp: ChatCompletion = self._openai.chat.completions.create(
            # 流模式
            stream=False,
            # 模型
            model=request.model.model_name,
            # 用户
            user=self._extract_end_user(request),
            # 消息
            messages=self._build_messages(request),
            # 响应结果数
            n=self._get_parameter('n', *params_array),
            # 核采样
            top_p=self._get_parameter('top_p', *params_array),
            # 采样温度
            temperature=self._get_parameter('temperature', *params_array),
            # 逻辑偏差
            logit_bias=self._get_parameter('logit_bias', *params_array),
            # 响应格式
            response_format=self._get_parameter('response_format', *params_array),
            # 存在惩罚
            presence_penalty=self._get_parameter('presence_penalty', *params_array),
            # 频率处罚
            frequency_penalty=self._get_parameter('frequency_penalty', *params_array),
        )
        # 提取令牌使用情况
        token_usage = self._extract_token_usage(resp)
        # 获取响应结果
        result = resp.choices[0]
        # 提取响应消息
        resp_msg = self._extract_messages(request.model, resp)
        # 检查消息是否完成
        if result.finish_reason != 'stop':
            # 若未完成，则尝试继续对话
            continued, next_msg, next_token_usage = self._try_continue_chat(
                result.finish_reason,
                request,
                resp_msg,
            )
            # 检查是否执行了继续对话
            if continued:
                # 合并响应结果
                resp_msg, token_usage = self._merge_result(
                    request.model,
                    resp_msg,
                    next_msg,
                    token_usage,
                    next_token_usage,
                )
        # 返回响应结果
        return (resp_msg,), token_usage
